#!/usr/bin/env python3
import rospy
import threading
import numpy as np
from std_msgs.msg import UInt32MultiArray
from sensor_msgs.msg import Image
from event_recorder.msg import CameraData
from rich import print
import os, sys, select, termios, tty
# 是否为debug模式
debugMode = True

# 定义发布者
publisherInitialized = False
nodePub = None

# 定义订阅者
subscriberInitialized = False
nodeSub = None

keyFlag = False
eventBuffer = []
imuBuffer = []

from cv_bridge import CvBridge
br = CvBridge()

termiosSettings = termios.tcgetattr(sys.stdin)

recordFileDir = os.path.join(os.environ["HOME"], "record_result")
simDir = os.path.join(os.environ["HOME"], "sim")

def getKey():
    global termiosSettings
    tty.setraw(sys.stdin.fileno())
    rlist, a, b = select.select([sys.stdin], [], [], 0.1)
    if rlist:
        key = sys.stdin.read(1)
    else:
        key = ""
    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, termiosSettings)
    return key

def readLock(dir):
    res = 0
    with open(dir, "r", encoding = "utf-8") as f:
        s = f.read()
        if s == "1":
            res = 1
    return res

def saveFile(savedEvents, savedIMUs, dir):
    os.makedirs(dir, exist_ok = True)
    fileName = "event_record" + str(rospy.Time.now())
    if savedEvents.shape[0]:
        savedEvents = savedEvents[np.argsort(savedEvents[:, 0])]
        # print(savedEvents)
        res = np.zeros((savedEvents.shape[0] * 2,), dtype = ">u4")
        res[::2] = savedEvents[:, 2] | (savedEvents[:, 3] << 1) | (savedEvents[:, 4] << 16)
        event_times = savedEvents[:, 0].astype("float64") + (savedEvents[:, 1].astype("float64") * 1e-9)
        # rospy.loginfo(np.min(event_times))
        # rospy.loginfo(np.max(event_times))
        event_times -= np.min(event_times)
        event_times *= 1e6
        event_times //= 1
        res[1::2] = event_times.astype(">u4")
        with open(os.path.join(dir, fileName + ".aedat"), "w+b") as file:
            file.write(res.tobytes())
        rospy.loginfo("%d event(s) saved.", savedEvents.shape[0])
    if savedIMUs.shape[0]:
        savedIMUs = savedIMUs[np.argsort(savedIMUs[:, 0])]
        np.savetxt(os.path.join(dir, fileName + ".csv"), savedIMUs, fmt = "%.8f", delimiter = ",")
    return fileName

"""
该节点的回调函数
@params:
    msg: EventVecInline 收到的（多个事件）消息
"""
def nodeSubCallback(msg: CameraData) -> None:
    global publisherInitialized, nodePub, br, imgPub, keyFlag, eventBuffer, imuBuffer, recordFileDir
    eventData = np.array(msg.events, dtype = "uint32")
    imuData = np.array(msg.imus, dtype = "float64")
    if eventData.shape[0]:
        eventData = eventData.reshape((eventData.shape[0] // 5, 5)) # [t_sec, t_nsec, p, y, x]
        if keyFlag:
            eventBuffer.append(eventData)
    if imuData.shape[0]:
        imuData = imuData.reshape((imuData.shape[0] // 4, 4)) # [t, wz, wy, wx]
        if keyFlag:
            imuBuffer.append(imuData)
    

"""
订阅者的线程函数，用于执行订阅者线程
"""
def subThreadFunc() -> None:
    global debugMode, subscriberInitialized, nodeSub
    nodeSub = rospy.Subscriber("/event_recorder/sub_event_res", CameraData, nodeSubCallback)  # 订阅事件消息
    subscriberInitialized = True
    rospy.spin()  # 开始循环监听消息


"""
发布者的线程函数，用于发布本结点的执行结果
"""
def pubThreadFunc() -> None:
    global debugMode, publisherInitialized, nodePub, imgPub, keyFlag, recordFileDir, simDir, eventBuffer, imuBuffer
    publisherInitialized = True
    rospy.loginfo("Files will be saved in %s", recordFileDir)
    rospy.loginfo("Press \"r\" to record.")
    while not rospy.is_shutdown():
        key = getKey()
        if key == "r":
            if not keyFlag:
                rospy.loginfo("Start recording...")
                rospy.loginfo("Press \"r\" to stop recording and save, \"q\" to stop recording without saving.")
                keyFlag = 1
            else:
                rospy.loginfo("Proessing data...")
                if len(eventBuffer):
                    savedEvents = np.vstack(eventBuffer)
                else:
                    savedEvents = np.array([])
                eventBuffer = []
                if len(imuBuffer):
                    savedIMUs = np.vstack(imuBuffer)
                else:
                    savedIMUs = np.array([])
                imuBuffer = []
                keyFlag = 0
                fileName = saveFile(savedEvents, savedIMUs, recordFileDir)
                rospy.loginfo("Record complete. Events are stored in %s.aedat" % os.path.join(recordFileDir, fileName))
                rospy.loginfo("Press \"r\" to record.")
        elif key == "q":
            keyFlag = 0
            eventBuffer = []
            imuBuffer = []
            rospy.loginfo("Recording stopped.")
            rospy.loginfo("Press \"r\" to record.")
        elif key == "\x03":
            break
        # lock = readLock(os.path.join(simDir, "record.lock"))
        # if lock and not keyFlag:
        #     keyFlag = 1
        #     rospy.loginfo("Record from simulator...")
        # if not lock and keyFlag:
        #     with open(os.path.join(simDir, "record_label.txt"), "r", encoding = "utf-8") as f:
        #         label = f.read()
        #     if len(eventBuffer):
        #         savedEvents = np.vstack(eventBuffer)
        #     else:
        #         savedEvents = np.array([])
        #     eventBuffer = []
        #     if len(imuBuffer):
        #         savedIMUs = np.vstack(imuBuffer)
        #     else:
        #         savedIMUs = np.array([])
        #     imuBuffer = []
        #     keyFlag = 0
        #     fileName = saveFile(savedEvents, savedIMUs, os.path.join(recordFileDir, label))
        #     rospy.loginfo("Record complete. Events are stored in %s.aedat" % os.path.join(recordFileDir, label, fileName))


if __name__ == "__main__":
    rospy.init_node("record", anonymous=True)  # 初始化ROS环境

    subT = threading.Thread(target=subThreadFunc)  # 定义订阅者线程
    pubT = threading.Thread(target=pubThreadFunc)  # 定义发布者线程

    subT.start()  # 启动订阅者线程
    pubT.start()  # 启动发布者线程

    subT.join()  # 加入订阅者线程
    pubT.join()  # 加入发布者线程
